import os
from PySide6.QtWidgets import (
    QApplication, QWidget, QLabel, QPushButton, QFileDialog, QHBoxLayout, QVBoxLayout,
    QGraphicsView, QGraphicsScene, QGraphicsPixmapItem, QMessageBox, QTreeWidget, QTreeWidgetItem,
    QSplitter, QFormLayout
)
from PySide6.QtGui import QPixmap, QImage, QPen
from PySide6.QtCore import Qt

import numpy as np
import cv2
import json

from .pipeline import rebuild_ui, imread_unicode


def np_to_qimage(img_bgr):
    if img_bgr is None:
        return QImage()
    rgb = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2RGB)
    h, w, ch = rgb.shape
    bytes_per_line = ch * w
    return QImage(rgb.data, w, h, bytes_per_line, QImage.Format_RGB888)


class UIRebuilderApp(QWidget):
    def __init__(self):
        super().__init__()
        self.setWindowTitle("UI Rebuilder")

        self.effect_path = os.path.join("资源", "效果图-设置界面.jpg")
        self.assets_dir = "资源"
        self.out_dir = "ui_out"

        self.scene = QGraphicsScene(self)
        self.view = QGraphicsView(self.scene)
        self.view.setDragMode(QGraphicsView.ScrollHandDrag)
        self.highlight_item = None
        self.result_data = None

        self.btn_effect = QPushButton("选择效果图")
        self.btn_assets = QPushButton("选择散图目录")
        self.btn_run = QPushButton("一键运行")

        self.lbl_effect = QLabel(self.effect_path)
        self.lbl_assets = QLabel(self.assets_dir)

        self.btn_effect.clicked.connect(self.choose_effect)
        self.btn_assets.clicked.connect(self.choose_assets)
        self.btn_run.clicked.connect(self.run_pipeline)

        top = QHBoxLayout()
        top.addWidget(self.btn_effect)
        top.addWidget(self.lbl_effect)
        top.addWidget(self.btn_assets)
        top.addWidget(self.lbl_assets)
        top.addWidget(self.btn_run)

        # Left: tree + props, Right: view
        self.tree = QTreeWidget()
        self.tree.setHeaderLabels(["名称", "类型"]) 
        self.tree.itemSelectionChanged.connect(self.on_tree_selection)

        self.prop_id = QLabel("-")
        self.prop_type = QLabel("-")
        self.prop_bounds = QLabel("-")
        self.prop_z = QLabel("-")
        self.prop_text = QLabel("-")
        self.prop_color = QLabel("-")
        self.prop_font = QLabel("-")

        prop_form = QFormLayout()
        prop_form.addRow("ID", self.prop_id)
        prop_form.addRow("类型", self.prop_type)
        prop_form.addRow("Bounds", self.prop_bounds)
        prop_form.addRow("Z", self.prop_z)
        prop_form.addRow("文本", self.prop_text)
        prop_form.addRow("颜色RGB", self.prop_color)
        prop_form.addRow("字号(px)", self.prop_font)

        left_panel = QWidget()
        left_layout = QVBoxLayout(left_panel)
        left_layout.addWidget(self.tree)
        left_layout.addLayout(prop_form)

        splitter = QSplitter()
        splitter.addWidget(left_panel)
        splitter.addWidget(self.view)
        splitter.setStretchFactor(0, 0)
        splitter.setStretchFactor(1, 1)

        layout = QVBoxLayout(self)
        layout.addLayout(top)
        layout.addWidget(splitter)
        self.setLayout(layout)

        # try initial preview
        self.refresh_preview()

    def choose_effect(self):
        path, _ = QFileDialog.getOpenFileName(self, "选择效果图", ".", "Images (*.png *.jpg *.jpeg *.bmp)")
        if path:
            self.effect_path = path
            self.lbl_effect.setText(path)
            self.refresh_preview()

    def choose_assets(self):
        path = QFileDialog.getExistingDirectory(self, "选择散图目录", ".")
        if path:
            self.assets_dir = path
            self.lbl_assets.setText(path)
            self.refresh_preview()

    def run_pipeline(self):
        try:
            result = rebuild_ui(self.effect_path, self.assets_dir, self.out_dir)
            QMessageBox.information(self, "完成", f"已输出:\n{result['json']}\n{result['preview']}")
            self.refresh_preview()
        except Exception as e:
            QMessageBox.critical(self, "错误", str(e))

    def refresh_preview(self):
        # load effect and overlay preview if exists
        effect = imread_unicode(self.effect_path)
        if effect is None:
            self.scene.clear()
            return
        qim = np_to_qimage(effect)
        pix = QPixmap.fromImage(qim)

        self.scene.clear()
        self.scene.addItem(QGraphicsPixmapItem(pix))

        preview_path = os.path.join(self.out_dir, 'preview.png')
        if os.path.exists(preview_path):
            overlay = imread_unicode(preview_path)
            if overlay is not None:
                qim2 = np_to_qimage(overlay)
                pix2 = QPixmap.fromImage(qim2)
                item2 = QGraphicsPixmapItem(pix2)
                item2.setOpacity(0.75)
                self.scene.addItem(item2)

        # Load tree/properties from result.json if present
        result_path = os.path.join(self.out_dir, 'result.json')
        self.result_data = None
        self.tree.clear()
        if os.path.exists(result_path):
            try:
                with open(result_path, 'r', encoding='utf-8') as f:
                    self.result_data = json.load(f)
                root = self.result_data.get('tree')
                if root:
                    self.populate_tree(None, root)
                    self.tree.expandToDepth(1)
            except Exception:
                pass

    def populate_tree(self, parent_item, node):
        name = node.get('name', node.get('type', 'node'))
        typ = node.get('type', '-')
        item = QTreeWidgetItem([str(name), str(typ)])
        item.setData(0, Qt.UserRole, node)
        if parent_item is None:
            self.tree.addTopLevelItem(item)
        else:
            parent_item.addChild(item)
        for child in node.get('children', []) or []:
            self.populate_tree(item, child)

    def on_tree_selection(self):
        items = self.tree.selectedItems()
        if not items:
            return
        node = items[0].data(0, Qt.UserRole)
        if not node:
            return
        self.update_props(node)
        self.highlight_bounds(node.get('bounds'))

    def update_props(self, node):
        self.prop_id.setText(str(node.get('id', '-')))
        self.prop_type.setText(str(node.get('type', '-')))
        b = node.get('bounds') or {}
        self.prop_bounds.setText(f"x={b.get('x','-')}, y={b.get('y','-')}, w={b.get('w','-')}, h={b.get('h','-')}")
        self.prop_z.setText(str(node.get('z_index', '-')))
        if node.get('type') == 'text':
            self.prop_text.setText(str(node.get('name', '')))
            style = node.get('style') or {}
            self.prop_color.setText(str(style.get('color_rgb', '-')))
            self.prop_font.setText(str(style.get('font_size_px', '-')))
        else:
            self.prop_text.setText('-')
            self.prop_color.setText('-')
            self.prop_font.setText('-')

    def highlight_bounds(self, bounds):
        if self.highlight_item is not None:
            self.scene.removeItem(self.highlight_item)
            self.highlight_item = None
        if not bounds:
            return
        x = int(bounds.get('x', 0))
        y = int(bounds.get('y', 0))
        w = int(bounds.get('w', 0))
        h = int(bounds.get('h', 0))
        pen = QPen(Qt.red)
        pen.setWidth(3)
        self.highlight_item = self.scene.addRect(x, y, w, h, pen)
        self.view.ensureVisible(self.highlight_item)


def run_gui():
    import sys
    app = QApplication(sys.argv)
    w = UIRebuilderApp()
    w.resize(1280, 800)
    w.show()
    sys.exit(app.exec())


if __name__ == "__main__":
    run_gui()


